#!/usr/bin/env python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import pathlib


class GrubEnvFile:

    _SIGNATURE = "# GRUB Environment Block"

    _WARNING = "# WARNING: Do not edit this file by tools other than grub-editenv!!!"

    _MIN_SIZE = 1024

    def __init__(self, bootDir):
        self._path = os.path.join(bootDir, "grub", "grubenv")
        self._bNeedSave, self._envDict = self._parse()

    @property
    def path(self):
        return self._path

    def is_saved(self):
        if self._bNeedSave:
            return False
        if not os.path.exists(self._path):
            return False
        # FIXME: should record and compare file mtime, so we can also know file is changed by others
        return True

    def has_env(self, name):
        return name in self._envDict

    def get_env(self, name):
        return self._envDict[name]

    def reset(self):
        if len(self._envDict) > 0:
            self._envDict.clear()
            self._bNeedSave = True

    def set_env(self, name, value):
        if self._envDict.get(name, None) != value:
            self._envDict[name] = value
            self._bNeedSave = True

    def unset_env(self, name):
        if name in self._envDict:
            del self._envDict[name]
            self._bNeedSave = True

    def save(self):
        if self.is_saved():
            return

        with open(self._path, "w") as f:
            count = 0
            if True:
                f.write(self._SIGNATURE + "\n")
                count += len(self._SIGNATURE) + 1
            if True:
                f.write(self._WARNING + "\n")
                count += len(self._WARNING) + 1
            for name, value in self._envDict.items():
                f.write("%s=%s\n" % (name, value))
                count += len(name) + 1 + len(value) + 1
            if count < self._MIN_SIZE:
                f.write("#" * (self._MIN_SIZE - count))
        self._bNeedSave = False

    def _parse(self):
        bNeedSave = False
        envDict = {}

        buf = None
        try:
            buf = pathlib.Path(self._path).read_text()
        except FileNotFoundError:
            return (True, envDict)

        if len(buf) < self._MIN_SIZE:
            # file is not legal (so it needs to be saved), but we still can parse it
            bNeedSave = True

        lineList = buf.split("\n")
        if len(lineList) == 0:
            return (True, envDict)
        if lineList[0] != self._SIGNATURE:
            return (True, envDict)

        for line in lineList[1:]:
            if line == "" or line.startswith("#"):
                continue
            m = re.fullmatch("(.*)=(.*)", line)
            if m is None:
                return (True, envDict)
            if m.group(1) in envDict:
                return (True, envDict)
            envDict[m.group(1)] = m.group(2)

        return (bNeedSave, envDict)


class GrubCfgFile:

    def __init__(self, bootDir):
        self._path = os.path.join(bootDir, "grub", "grub.cfg")
        try:
            self._content = pathlib.Path(self._path).read_text()
            self._bNeedSave = False
        except FileNotFoundError:
            self._content = ""
            self._bNeedSave = True

    @property
    def path(self):
        return self._path

    def is_saved(self):
        if self._bNeedSave:
            return False
        if not os.path.exists(self._path):
            return False
        # FIXME: should record and compare file mtime, so we can also know file is changed by others
        return True

    def get_content(self):
        return self._content

    def set_content(self, content):
        if self._content != content:
            self._content = content
            self._bNeedSave = True

    def save(self):
        if self.is_saved():
            return

        with open(self._path, "w") as f:
            f.write(self._content)
        self._bNeedSave = False
